refactor(attn): let the bucketer define the batched decode kernel's buckets - #792
refactor(attn): let the bucketer define the batched decode kernel's buckets#792jvlunteren wants to merge 2 commits into
Conversation
|
👋 Hi! Thank you for contributing. We also recommend installing prek and configuring it to check your code before every local commit. |
f58a98b to
2019878
Compare
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
2019878 to
620c48d
Compare
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
620c48d to
5084ae2
Compare
sducouedic
left a comment
sducouedic
left a comment
There was a problem hiding this comment.
Look good and clean! just a two small comments
| # The buckets are sized from max_model_len; a small KV allocation | ||
| # cannot host that many distinct pages to gather. | ||
| continue |
There was a problem hiding this comment.
should we print a warning here, or even an assert? I think vllm upstream enforces --num-gpu-blocks-override to serve at least one max-model-len request
| reachable = [n for n in self._num_seqs_buckets if n >= _MIN_BATCHED_SEQS] | ||
| for num_seqs in sorted(reachable, reverse=True): |
There was a problem hiding this comment.
as in the other PR, is there anything preventing us to set _num_seqs_buckets with values already >= _MIN_BATCHED_SEQS
|
Superseded by #810. Most of this has landed independently since the merge-base ( The idea that's still un-done is the one this PR opens with: the bucketer should own the ladders. #810 does that, with you as co-author. It also fixes a related divergence I hit while extracting it — the builder's The builder dispatches onto Leaving this open for you to close, in case you see something in |
2 participants
Description
Makes
SpyreAttnBucketerthe single source of truth for attention bucket sizes, and has warmup pre-compile the batched decode kernel as well as the per-sequence one. Per review feedback on #772.The metadata builder computed its own seqs and blocks buckets in parallel with the bucketer computing the same thing, so the set
build()dispatched onto and the set warmup recorded agreed only by coincidence. It now reads both from the bucketer. The batched kernel keys on a different tuple from the per-sequence one, so it gets its own bucket type and enumeration, and record_graphs compiles both. Previously the batched kernel compiled lazily in the serving path, the stall warmup exists to prevent. Buckets dispatch cannot reach are pruned.Rebased onto #789, which routes query and output through fixed-size staging buffers:
needs_gathernow comes from the sequence counts rather than the query width, which staging made constant, and the recorder records on those buffers instead of fresh tensors.Stacked on #791, since it uses the new names throughout.
Related Issues
#772 #789 #791
Test Plan
pytest tests/runtime/test_spyre_attn_bucketer.py: 69 passedpytest tests/attention/test_spyre_attn.py --attn-shards=4 --attn-shard-id=0: 46 passedChecklist
bash format.sh)Signed-off-by:line (DCO compliance)